{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "create by 2018-05-18\n",
    "\n",
    "@author: Shiyipaisizuo\n",
    "\"\"\"\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "# 定义文本框和箭头格式\n",
    "decisionnode = dict(boxstyle='sawtooth', fc='0.8')\n",
    "leafnode = dict(boxstyle='round4', fc='0.8')\n",
    "arrow_args = dict(arrowstyle='<-')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 绘制带箭头的注释\n",
    "def plot_node(nodetxt, centerpt, parentpt, nodetype):\n",
    "    create_plot.ax1.annotate(nodetxt, xy=parentpt, xycoords='axes fraction',\n",
    "                            xytext=centerpt, textcoords='axes fraction',\n",
    "                            va=\"center\", ha=\"center\", bbox=nodetype, arrowprops=arrow_args)\n",
    "\n",
    "# 获取叶子节点的数目\n",
    "def get_num_leafs(mytree):\n",
    "    numleafs = 0\n",
    "    firststr = mytree.keys()[0]\n",
    "    seconddict = mytree[firststr]\n",
    "\n",
    "    # 测试节点数据是否为字典\n",
    "    for key in seconddict.keys():\n",
    "        if type(seconddict[key]).__name__ == 'dict':\n",
    "            numleafs += get_num_leafs(seconddict[key])\n",
    "        else:\n",
    "            numleafs += 1\n",
    "    return numleafs\n",
    "\n",
    "\n",
    "# 获取决策树的层数\n",
    "def get_tree_depth(mytree):\n",
    "    maxdepth = 0\n",
    "    firststr = mytree.keys()[0]\n",
    "    seconddict = mytree[firststr]\n",
    "\n",
    "    # 测试数据是否为字典\n",
    "    for key in seconddict.keys():\n",
    "        if type(seconddict[key]).__name__ == 'dict':\n",
    "            thisdepth = 1 + get_tree_depth(seconddict[key])\n",
    "        else:\n",
    "            thisdepth = 1\n",
    "        if thisdepth > maxdepth: maxdepth = thisdepth\n",
    "    return maxdepth\n",
    "\n",
    "\n",
    "# 预先储存树的信息，避免每次都要从数据中创建树的麻烦\n",
    "def retrieve_tree(i):\n",
    "    listoftrees = [{'no surfacing': {0: 'no', 1: {'flippers': {0: 'no', 1: 'yes'}}}},\n",
    "                   {'no surfacing': {0: 'no', 1: {'flippers': {0: {'head': {0: 'no', 1: 'yes'}}, 1: 'no'}}}}\n",
    "                   ]\n",
    "    return listoftrees[i]\n",
    "\n",
    "\n",
    "# 在父节点中填充文本的信息\n",
    "def plot_mid_text(cntrpt, parentpt, txtstring):\n",
    "    xmid = (parentpt[0] - cntrpt[0]) / 2.0 + cntrpt[0]\n",
    "    ymid = (parentpt[1] - cntrpt[1]) / 2.0 + cntrpt[1]\n",
    "    create_plot.ax1.text(xmid, ymid, txtstring, va=\"center\", ha=\"center\", rotation=30)\n",
    "\n",
    "\n",
    "# 绘制树\n",
    "def plot_tree(mytree, parentpt, nodetxt):\n",
    "\n",
    "    # 计算树的宽和高\n",
    "    numleafs = get_num_leafs(mytree)\n",
    "    depth = get_tree_depth(mytree)\n",
    "    firststr = mytree.keys()[0]\n",
    "    cntrpt = (plot_tree.xoff + (1.0 + float(numleafs)) / 2.0 / plot_tree.totalw, plot_tree.yoff)\n",
    "    plot_mid_text(cntrpt, parentpt, nodetxt)\n",
    "    plot_node(firststr, cntrpt, parentpt, decisionnode)\n",
    "    seconddict = mytree[firststr]\n",
    "    plot_tree.yoff = plot_tree.yoff - 1.0 / plot_tree.totald\n",
    "\n",
    "    # 测试数据是否为字典\n",
    "    for key in seconddict.keys():\n",
    "        if type(seconddict[key]).__name__ == 'dict':\n",
    "            plot_tree(seconddict[key], cntrpt, str(key))\n",
    "\n",
    "        else:  # 它是叶子节点打印叶子节点\n",
    "            plot_tree.xoff = plot_tree.xoff + 1.0 / plot_tree.totalw\n",
    "            plot_node(seconddict[key], (plot_tree.xoff, plot_tree.yoff), cntrpt, leafnode)\n",
    "            plot_mid_text((plot_tree.xoff, plot_tree.yoff), cntrpt, str(key))\n",
    "    plot_tree.yoff = plot_tree.yoff + 1.0 / plot_tree.totald\n",
    "\n",
    "\n",
    "# 如果你得到了一个dictonary，你知道它是一棵树，第一个元素将是另一个法令。\n",
    "def create_plot(intree):\n",
    "    fig = plt.figure(1, facecolor='white')\n",
    "    fig.clf()\n",
    "    axprops = dict(xticks=[], yticks=[])\n",
    "    create_plot.ax1 = plt.subplot(111, frameon=False, **axprops)  # no ticks\n",
    "    plot_tree.totalw = float(get_num_leafs(intree))\n",
    "    plot_tree.totald = float(get_tree_depth(intree))\n",
    "    plot_tree.xoff = -0.5 / plot_tree.totalw;\n",
    "    plot_tree.yoff = 1.0;\n",
    "    plot_tree(intree, (0.5, 1.0), '')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
